from ultralytics import YOLO
import torch
from torch import nn
if __name__ == '__main__':

    model = YOLO(model=r'D:\Traffic_intersection\Helmet_detection\ultralytics-main\ultralytics\cfg\models\11\yolo11.yaml')
    model.load('yolo11n.pt')
    # model.train(data=r'D:\Traffic_intersection\Helmet_detection\ultralytics-main\mydata.yaml',
    #             imgsz=640,
    #             epochs=50,
    #             batch=4,
    #             device='0'
    #             )
    # 计算类别权重
    model.train(data=r'D:\Traffic_intersection\Helmet_detection\ultralytics-main\mydata.yaml',
                imgsz=640,
                epochs=70,
                batch=8,
                device=0,
                flipud=0.5,
                fliplr=0.5,
                degrees=10,
                mosaic=1.0,  # 启用 Mosaic 数据增强
                mixup=0.5,  # 启用 Mixup 数据增强
                hsv_h=0.015, # 色调增强，范围为 ±0.015
                hsv_s=0.7, # 饱和度增强，范围为 ±0.7
                hsv_v=0.4, # 亮度增强，范围为 ±0.4
                scale=0.5,  # 随机缩放，范围为 ±50%
                lr0=0.001,
                lrf=0.1,
                optimizer='SGD'
                )





